/*!
*  Copyright (c) 2015 by Joshua Zhang
* \file detector.hpp
* \brief ssd detection module impl
*/
#include "zupply.hpp"
#include "ssddetector.hpp"
#include <iostream>
#include <sstream>
#include <streambuf>
#include <cassert>
#include <cstdlib>
#include <ctime>

using namespace zz;

SSDDetector::SSDDetector(std::string model_prefix, int epoch, int width,
	int height, float mean_r, float mean_g, float mean_b,
	int device_type, int device_id) {
	auto logger = log::get_logger("default");
	if (epoch < 0 || epoch > 9999) {
		logger->error("Invalid epoch number: ") << epoch;
		exit(-1);
	}
	std::string model_file = model_prefix + "-" + fmt::int_to_zero_pad_str(epoch, 4) + ".params";
	if (!os::is_file(model_file)) {
		std::cerr << "Model file: " << model_file << " does not exist" << std::endl;
		exit(-1);
	}
	std::string json_file = model_prefix + "-symbol.json";
	if (!os::is_file(json_file)) {
		std::cerr << "JSON file: " << model_file << " does not exist" << std::endl;
		exit(-1);
	}
	if (width < 1 || height < 1) {
		std::cerr << "Invalid width or height: " << width << "," << height << std::endl;
		exit(-1);
	}
	width_ = width;
	height_ = height;
	const char *input_name = "data";
	const char *input_keys[1];
	input_keys[0] = input_name;
	const mx_uint input_shape_indptr[] = { 0, 4 };
	const mx_uint input_shape_data[] = { 1, 3, static_cast<mx_uint>(width_), static_cast<mx_uint>(height_) };
	mean_r_ = mean_r;
	mean_g_ = mean_g;
	mean_b_ = mean_b;

	// load model
	std::ifstream param_file(model_file, std::ios::binary | std::ios::ate);
	if (!param_file.is_open()) {
		std::cerr << "Unable to open model file: " << model_file << std::endl;
		exit(-1);
	}
	std::streamsize size = param_file.tellg();
	param_file.seekg(0, std::ios::beg);
	buffer_.resize(size);

	std::ifstream json_handle(json_file, std::ios::ate);
	std::string json_buffer;
	json_buffer.reserve(json_handle.tellg());
	json_handle.seekg(0, std::ios::beg);
	json_buffer.assign((std::istreambuf_iterator<char>(json_handle)), std::istreambuf_iterator<char>());
	if (json_buffer.size() < 1) {
		std::cerr << "invalid json file: " << json_file << std::endl;
		exit(-1);
	}

	if (param_file.read(buffer_.data(), size)) {
		MXPredCreate(json_buffer.c_str(), buffer_.data(), static_cast<int>(size),
			device_type, device_id, 1, input_keys, input_shape_indptr, input_shape_data,
			&predictor_);
	}
	else {
		std::cerr << "Unable to read model file: " << model_file << std::endl;
		exit(-1);
	}
}

std::vector<float> SSDDetector::detect(std::string in_img) {
	cv::Mat img = cv::imread(in_img);
	return detect(img);
}

void SSDDetector::detect(const cv::Mat& img, std::vector<cv::Point2f>& centers, std::vector<cv::Rect>& rects)
{
	auto&& outputs = detect(img);

	int width = img.cols;
	int height = img.rows;
	for (int i = 0; i < outputs.size(); i += 6) {
		if (outputs[i] < 0) continue;  // not an object
		int id = static_cast<int>(outputs[i] + 0.5);
		float score = outputs[i + 1];
		if (score < 0.5) continue;
		float xmin = outputs[i + 2];
		float ymin = outputs[i + 3];
		float xmax = outputs[i + 4];
		float ymax = outputs[i + 5];

		centers.emplace_back((xmin + xmax) / 2.*width, (ymin + ymax) / 2.*height);
		rects.emplace_back(cv::Point(xmin*width, ymin*height), cv::Point(xmax*width, ymax*height));
	}
}

std::vector<float> SSDDetector::detect(const cv::Mat& image) {
	auto logger = log::get_logger("default");
	if (image.empty()) {
		std::cerr << "Image is empty" << std::endl;
		exit(-1);
	}
	if (image.channels() != 3) {
		std::cerr << "RGB image required" << std::endl;
		exit(-1);
	}
	// prepare data
	cv::Mat resize_img;
	cv::resize(image, resize_img, cv::Size(width_, height_));
	int size = resize_img.channels() * width_ * height_;
	cv::Mat image_float;
	resize_img.convertTo(image_float, CV_32FC3);
	cv::Mat bgr[3];
	cv::split(image_float, bgr);
	bgr[0] = bgr[0] - mean_b_;
	bgr[1] = bgr[1] - mean_g_;
	bgr[2] = bgr[2] - mean_r_;
	cv::Mat stack = bgr[2];
	stack.push_back(bgr[1]);
	stack.push_back(bgr[0]);
	// use model to forward
	mx_uint *shape = NULL;
	mx_uint shape_len = 0;
	MXPredSetInput(predictor_, "data", (float*)stack.data, static_cast<mx_uint>(size));
	time::Timer timer;
	MXPredForward(predictor_);
	MXPredGetOutputShape(predictor_, 0, &shape, &shape_len);
	mx_uint tt_size = 1;
	for (mx_uint i = 0; i < shape_len; ++i) {
		tt_size *= shape[i];
	}
	assert(tt_size % 6 == 0);
	std::vector<float> outputs(tt_size);
	MXPredGetOutput(predictor_, 0, outputs.data(), tt_size);
	//logger->info("Detect forward elapsed time: ") << timer.to_string();
	return outputs;
}
